from functools import cached_property
from typing import Iterable, TypeVar

import numpy as np
from numpy.typing import ArrayLike

from .chained import ChainedDatasets

AnySample = TypeVar('AnySample')


class KnownSizeChainedDatasets(ChainedDatasets[AnySample]):
    """ A dataset that produces samples from a (manageable sized) sequence of sub-datasets in order. """

    def __init__(self,
        children: Iterable[Iterable[AnySample]],
        children_samples_n: ArrayLike | None = None,
    ):
        super().__init__(children)
        self._children_samples_n = np.asarray(children_samples_n, dtype=np.int64) \
            if children_samples_n is not None else None

    @property
    def children_samples_n(self) -> np.ndarray:
        if self._children_samples_n is None:
            children_samples_n = [len(child) for child in self._children]
            children_samples_n = np.asarray(children_samples_n, dtype=np.int64)
            self._children_samples_n = children_samples_n
        return self._children_samples_n

    @cached_property
    def children_tail(self) -> np.ndarray:
        return np.cumsum(self.children_samples_n)

    def locate(self, index: int, /) -> tuple[int, int]:
        if len(self.children_tail) == 0:
            raise IndexError
        if not 0 <= index < self.children_tail[-1]:
            raise IndexError
        child_index = int(np.searchsorted(self.children_tail, index, side='right'))
        child_head = 0 if child_index == 0 else self.children_tail[child_index - 1]
        sample_index = index - child_head
        return child_index, sample_index

    def __len__(self) -> int:
        if len(self.children_tail) == 0:
            return 0
        return int(self.children_tail[-1])

    def __getitem__(self, index: int, /) -> AnySample:
        child_index, sample_index = self.locate(index)
        # noinspection PyUnresolvedReferences
        return self._children[child_index][sample_index]
